# https://gitee.com/yueyinqiu5990/tj12413601/blob/master/assignment4/question3/main.py
import matplotlib.pyplot
import numpy

import electrostatic_solver


def _field_base(size: int = 51):
    field = numpy.empty([size, size], dtype=numpy.object_)
    field.fill(None)

    field[0].fill(0.)
    field[-1].fill(0.)
    field[:, 0].fill(0.)
    field[:, -1].fill(0.)

    field[size // 2, size // 2] = 1.

    return field, size


def _solve_and_plot(subplot: int, title: str, field):
    field = electrostatic_solver.solve_dirichlet(field)

    matplotlib.pyplot.subplot(subplot)
    matplotlib.pyplot.jet()

    row_count, column_count = field.shape
    x = numpy.linspace(-1, 1, row_count)
    y = numpy.linspace(-1, 1, column_count)
    x, y = numpy.meshgrid(x, y)
    matplotlib.pyplot.contourf(x, y, field, levels=100, extend="both")

    matplotlib.pyplot.title(f"{title}")
    matplotlib.pyplot.xlabel("$x$")
    matplotlib.pyplot.ylabel("$y$")
    matplotlib.pyplot.colorbar()


def _main():
    print(f"Solving problem (1)...")
    field, size = _field_base()
    _solve_and_plot(221, "Problem (1)", field)

    print(f"Solving problem (2)...")
    field, size = _field_base()
    theta = numpy.linspace(0, 2.1 * numpy.pi, 100)
    x = numpy.rint(8 * numpy.cos(theta) + size // 2).astype(numpy.int_)
    y = numpy.rint(8 * numpy.sin(theta) + size // 2).astype(numpy.int_)
    for x, y in zip(x, y):
        field[x, y] = 0.
    _solve_and_plot(222, "Problem (2)", field)

    print(f"Solving problem (3)...")
    field, size = _field_base()
    theta = numpy.linspace(0, 2.1 * numpy.pi, 100)
    x = numpy.rint(8 * numpy.cos(theta) + size * 0.25).astype(numpy.int_)
    y = numpy.rint(8 * numpy.sin(theta) + size * 0.25).astype(numpy.int_)
    for x, y in zip(x, y):
        field[x, y] = 0.
    _solve_and_plot(223, "Problem (3)", field)

    print(f"Solving problem (4)...")
    field, size = _field_base()
    theta = numpy.linspace(0, 2.1 * numpy.pi, 10000)
    rho = 10 * (1 - numpy.sin(theta))
    # 这里 + 3 以避免爱心覆盖掉电荷
    x = numpy.rint(rho * numpy.sin(theta) + size // 2 + 3).astype(numpy.int_)
    y = numpy.rint(rho * numpy.cos(theta) + size // 2).astype(numpy.int_)
    for x, y in zip(x, y):
        field[x, y] = 0.
    _solve_and_plot(224, "Problem (4)", field)

    matplotlib.pyplot.subplots_adjust(
        left=0.229, bottom=0.11, right=0.752, top=0.905,
        wspace=0.42, hspace=0.298)
    matplotlib.pyplot.show()


if __name__ == "__main__":
    _main()
